package middleware

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/zeromicro/go-zero/core/codec"
	"github.com/zeromicro/go-zero/core/iox"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
	"io"
	"mall-backend/api/internal/config"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	CodeClientErr = 1
)

type SignatureMiddleWare struct {
	cfg        config.Config
	decrypters map[string]codec.RsaDecrypter
}

type ErrData struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

func ResponseError(code int, msg string) *ErrData {
	return &ErrData{
		Code:    code,
		Message: msg,
	}
}

const (
	requestUriHeader = "X-Request-Uri"
	signatureField   = "signature"
	timeField        = "time"
	contentSecurity  = "X-Content-Security"
)

var (
	// ErrInvalidContentType is an error that indicates invalid content type.
	ErrInvalidContentType = errors.New("invalid content type")
	// ErrInvalidHeader is an error that indicates invalid X-Content-Security header.
	ErrInvalidHeader = errors.New("invalid X-Content-Security header")
	// ErrInvalidKey is an error that indicates invalid key.
	ErrInvalidKey = errors.New("invalid key")
	// ErrInvalidPublicKey is an error that indicates invalid public key.
	ErrInvalidPublicKey = errors.New("invalid public key")
	// ErrInvalidSecret is an error that indicates invalid secret.
	ErrInvalidSecret = errors.New("invalid secret")
)

type ContentSecurityHeader struct {
	Key         []byte
	Timestamp   string
	ContentType int
	Signature   string
}

func NewSignatureMiddleWare(c config.Config) *SignatureMiddleWare {
	middleware := &SignatureMiddleWare{
		cfg: c,
	}
	decrypters, err := middleware.parseSignatureConf()
	if err != nil {
		logx.Error("NewSignatureMiddleWare middleware.parseSignatureConf err:", err)
		panic(err)
	}
	middleware.decrypters = decrypters
	return middleware
}

// Handle 签名校验处理
func (m *SignatureMiddleWare) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 鉴权关闭
		if !m.cfg.SignConf.Strict {
			next(w, r)
			return
		}
		if r.Body != nil {
			buf, err := io.ReadAll(r.Body)
			if err != nil {
				logx.Error("SignatureMiddleWare io.ReadAll err:", err)
				httpx.WriteJson(w, http.StatusBadRequest, ResponseError(CodeClientErr, "http body err"))
				return
			}
			if len(buf) > 0 {
				jsonBody := make(map[string]interface{})
				err = json.Unmarshal(buf, &jsonBody)
				if err != nil {
					logx.Error("SignatureMiddleWare json.Unmarshal err:", err)
					httpx.WriteJson(w, http.StatusBadRequest, ResponseError(CodeClientErr, "http body not json"))
					return
				}
				buf, err = json.Marshal(&jsonBody)
				if err != nil {
					logx.Error("SignatureMiddleWare json.Marshal err:", err)
					httpx.WriteJson(w, http.StatusBadRequest, ResponseError(CodeClientErr, "http body not json"))
					return
				}
			}
			r.Body = io.NopCloser(bytes.NewBuffer(buf))
		}
		header, err := m.parseSignatureHeader(r)
		if err != nil {
			logx.Error("SignatureMiddleWare parseSignatureHeader err:", err)
			httpx.WriteJson(w, http.StatusBadRequest, ResponseError(CodeClientErr, err.Error()))
			return
		}
		code := m.verifySignature(r, header)
		if code != httpx.CodeSignaturePass {
			logx.Errorf("Signature verification failed, X-Content-Security: %s",
				r.Header.Get(contentSecurity))
			httpx.WriteJson(w, http.StatusBadRequest, ResponseError(CodeClientErr, "signature err"))
			return
		}
		// 处理请求
		next(w, r)
	}
}

func (m *SignatureMiddleWare) verifySignature(r *http.Request, securityHeader *ContentSecurityHeader) int {
	seconds, err := strconv.ParseInt(securityHeader.Timestamp, 10, 64)
	if err != nil {
		return httpx.CodeSignatureInvalidHeader
	}
	now := time.Now().Unix()
	toleranceSeconds := 360 * 24 * 3600 * m.cfg.SignConf.Expiry
	if seconds+toleranceSeconds < now || now+toleranceSeconds < seconds {
		return httpx.CodeSignatureWrongTime
	}
	reqPath, reqQuery := m.getPathQuery(r)
	signContent := strings.Join([]string{
		securityHeader.Timestamp,
		r.Method,
		reqPath,
		reqQuery,
		m.computeBodySignature(r),
	}, "\n")

	actualSignature := codec.HmacBase64(securityHeader.Key, signContent)
	logx.Infof("sign content timestamp:%s method:%s reqPath:%s reqQuery:%s result:%s", securityHeader.Timestamp, r.Method, reqPath, reqQuery, actualSignature)
	if securityHeader.Signature == actualSignature {
		return httpx.CodeSignaturePass
	}
	logx.Errorf("signature different, expect: %s, actual: %s",
		securityHeader.Signature, actualSignature)
	return httpx.CodeSignatureInvalidToken
}

func (m *SignatureMiddleWare) parseSignatureConf() (map[string]codec.RsaDecrypter, error) {
	decrypters := make(map[string]codec.RsaDecrypter)
	decrypter, err := codec.NewRsaDecrypter(m.cfg.SignConf.PrivateKey)
	if err != nil {
		return nil, err
	}
	decrypters[m.cfg.SignConf.Fingerprint] = decrypter
	return decrypters, nil
}

func (m *SignatureMiddleWare) computeBodySignature(r *http.Request) string {
	var dup io.ReadCloser
	r.Body, dup = iox.DupReadCloser(r.Body)
	sha := sha256.New()
	io.Copy(sha, r.Body)
	r.Body = dup
	signBody := fmt.Sprintf("%x", sha.Sum(nil))
	logx.Infof("signBody:%s", signBody)
	return signBody
}

func (m *SignatureMiddleWare) getPathQuery(r *http.Request) (string, string) {
	requestUri := r.Header.Get(requestUriHeader)
	if len(requestUri) == 0 {
		return r.URL.Path, r.URL.RawQuery
	}

	uri, err := url.Parse(requestUri)
	if err != nil {
		return r.URL.Path, r.URL.RawQuery
	}

	return uri.Path, uri.RawQuery
}

func (m *SignatureMiddleWare) parseSignatureHeader(r *http.Request) (
	*ContentSecurityHeader, error) {
	contentSecurity := r.Header.Get(httpx.ContentSecurity)
	attrs := httpx.ParseHeader(contentSecurity)
	fingerprint := attrs[httpx.KeyField]
	secret := attrs[httpx.SecretField]
	signature := attrs[signatureField]

	if len(fingerprint) == 0 || len(secret) == 0 || len(signature) == 0 {
		return nil, ErrInvalidHeader
	}

	decrypter, ok := m.decrypters[fingerprint]
	if !ok {
		return nil, ErrInvalidPublicKey
	}

	decryptedSecret, err := decrypter.DecryptBase64(secret)
	if err != nil {
		return nil, ErrInvalidSecret
	}

	attrs = httpx.ParseHeader(string(decryptedSecret))
	base64Key := attrs[httpx.KeyField]
	timestamp := attrs[timeField]
	contentType := attrs[httpx.TypeField]

	key, err := base64.StdEncoding.DecodeString(base64Key)
	if err != nil {
		return nil, ErrInvalidKey
	}

	cType, err := strconv.Atoi(contentType)
	if err != nil {
		return nil, ErrInvalidContentType
	}
	logx.Infof("key:%s, timestamp:%s ctype:%d signature:%s", string(key), timestamp, cType, signature)
	return &ContentSecurityHeader{
		Key:         key,
		Timestamp:   timestamp,
		ContentType: cType,
		Signature:   signature,
	}, nil
}
